# Import required libraries
import gym_tanks
import gymnasium as gym
from stable_baselines3 import PPO
from multiprocessing import freeze_support

# 解决 macOS 上的多重库加载问题
import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'

if __name__ == "__main__":
    # 避免多进程问题
    freeze_support()

    # 创建并初始化游戏环境
    env = gym.make('gym_tanks/tanks-v0')
    env.reset()

    # 设置每个训练步骤的时间步长,收集训练的数据,在收集了10000次后才能进行一次策略更新
    TIMESTEPS = 10000

    # 设置模型保存的文件夹和路径
    models_dir = "models/PPO/Test_Bot"
    model_path = f"{models_dir}/model_2064384_steps.zip" # 在这里指定要加载的权重文件

    # 加载预训练的 PPO 模型
    model = PPO.load(model_path, env = env, n_steps = TIMESTEPS)

    # 向量化环境以提高性能
    env = model.get_env()

    # 设置测试回合数
    episodes = 10

    # 运行测试回合
    for ep in range(episodes):
        # 重置环境,获取初始观察
        obs = env.reset()
        done = False
        
        # 运行单个回合直到结束
        while not done:
            # 渲染游戏画面
            env.render()
            
            # 使用模型预测动作, deterministic可以设置是否是确定性策略还是随机策略
            action, _state = model.predict(obs, deterministic =True)            
            # 执行动作,获取环境反馈
            obs, reward, done, info = env.step(action)

            # # 打印当前奖励值
            # print(reward)

    # 关闭环境
    env.close()
